# -*- coding: utf-8 -*-
# @Author  : zlh
# @Time    : 2025/3/26
# @Function: 标准化GFS NetCDF文件变量，使新旧文件变量保持一致

import os
import xarray as xr
import numpy as np
from loguru import logger
import pandas as pd
from pathlib import Path
import glob
from core.download_gfs import grb2_to_nc 
def standardize_gfs_variables(new_file, reference_file, output_file=None):
    """
    将新的GFS NetCDF文件中的变量与参考文件保持一致
    
    参数:
        new_file: 需要标准化的新GFS文件路径
        reference_file: 参考GFS文件路径
        output_file: 输出文件路径，如果为None则覆盖原文件
    
    返回:
        标准化后的文件路径
    """
    logger.info(f"开始标准化文件: {new_file}")
    
    if output_file is None:
        output_file = new_file
    
    try:
        # 读取数据集，显式指定engine
        ds_new = xr.open_dataset(new_file, engine='netcdf4')
        ds_ref = xr.open_dataset(reference_file, engine='netcdf4')
        
        # 统一坐标名称
        coord_map = {}
        if 'latitude' in ds_new.dims:
            coord_map['latitude'] = 'lat'
        if 'longitude' in ds_new.dims:
            coord_map['longitude'] = 'lon'
        ds_new = ds_new.rename(coord_map)
        
        # 提取时间信息
        file_name = Path(new_file).name
        date_str = file_name.split('.')[2][:8]
        hour_str = file_name.split('.')[2][8:10]
        new_time = pd.to_datetime(f"{date_str}{hour_str}", format='%Y%m%d%H')
        
        # 处理降水数据
        if 'tp' in ds_new:
            data = ds_new.tp.values
            expanded_data = np.expand_dims(data, axis=0)
            
            # 创建新数据集，使用.data或.values提取数组数据
            ds_output = xr.Dataset(
                data_vars={
                    'A_PCP_L1_Accum_1': (['time', 'lat', 'lon'], expanded_data)
                },
                coords={
                    'time': ('time', [new_time]),
                    'lat': ('lat', ds_new.lat.values),  # 使用.values获取纯数组数据
                    'lon': ('lon', ds_new.lon.values)   # 使用.values获取纯数组数据
                }
            )
            
            # 复制降水变量的属性
            if 'A_PCP_L1_Accum_1' in ds_ref:
                ds_output['A_PCP_L1_Accum_1'].attrs.update(ds_ref['A_PCP_L1_Accum_1'].attrs)
            
            # 设置编码
            encoding = {
                'A_PCP_L1_Accum_1': {'zlib': True, 'complevel': 5},
                'time': {'dtype': 'int64'},
                'lat': {'dtype': 'float32'},
                'lon': {'dtype': 'float32'}
            }
            
            # 保存文件
            ds_output.to_netcdf(output_file, encoding=encoding)
            ds_output.close()
            
            # 关闭数据集
            ds_new.close()
            ds_ref.close()
            
            return output_file
            
    except Exception as e:
        logger.error(f"标准化文件时出错: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return None

def process_directory(input_dir, reference_file, output_dir):
    """
    处理输入目录中的所有NetCDF文件
    
    参数:
        input_dir: 输入目录路径
        reference_file: 参考文件路径
        output_dir: 输出目录路径
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取输入目录中的所有grib2文件
    input_files = glob.glob(os.path.join(input_dir, "gfs.0p25.*.grib2"))
    
    # 先将所有pgrb2文件转换为nc文件
    for pgrb2_file in input_files:
        # 构造nc文件路径
        nc_file = pgrb2_file + '.nc'
        
        try:
            # 转换pgrb2到nc
            # 如果nc文件已存在，则跳过转换
            if os.path.exists(nc_file):
                #logger.info(f"文件 {nc_file} 已存在，跳过转换")
                continue
            
            grb2_to_nc(pgrb2_file, nc_file)
            logger.info(f"成功转换文件: {pgrb2_file} -> {nc_file}")
            
           
        except Exception as e:
            logger.error(f"处理文件 {pgrb2_file} 时出错: {e}")
            continue
    
    # 获取转换后的所有nc文件进行处理
    input_files = glob.glob(os.path.join(input_dir, "*.nc"))
    
    logger.info(f"在目录 {input_dir} 中找到 {len(input_files)} 个文件")
    
    # 统计成功和失败的文件
    success_files = []
    failed_files = []
    
    # 处理每个文件
    for input_file in input_files:
        # 构建输出文件路径
        file_name = os.path.basename(input_file)
        output_file = os.path.join(output_dir, file_name)
        
        # 检查并删除已存在的输出文件
        if os.path.exists(output_file):
            os.remove(output_file)
            #logger.info(f"删除已存在的输出文件: {output_file}")
        
        # 标准化文件
        result = standardize_gfs_variables(input_file, reference_file, output_file)
        
        if result:
            success_files.append(file_name)
        else:
            failed_files.append(file_name)
    
    # 输出统计信息
    logger.info(f"目录 {input_dir} 处理完成")
    logger.info(f"成功处理: {len(success_files)}/{len(input_files)} 个文件")
    
    if failed_files:
        logger.warning(f"处理失败: {len(failed_files)} 个文件")
        for failed_file in failed_files:
            logger.warning(f"  - {failed_file}")
    
    return len(success_files), len(failed_files), failed_files
        
if __name__ == "__main__":
    # 配置日志
    # logger.add("/home/zlh/rainfall/logs/gfs_standardizer.log", rotation="10 MB")
    
    # 设置目录和参考文件
    input_dir = "/home/zlh/rainfall/data/GFS_minio/48"
    reference_file = "/home/zlh/rainfall/data/GFS/24/gfs.0p25.2021010100.f024.grib2.nc"
    output_dir = "/home/zlh/rainfall/data/stard_gfs/48"
    
    # 处理整个目录
    success_count, failed_count, failed_files = process_directory(input_dir, reference_file, output_dir)
    
    # 输出总结
    print(f"\n处理总结:")
    #print(f"总文件数: {success_count + failed_count}")
    print(f"成功处理: {success_count}")
    print(f"处理失败: {failed_count}")
    
    if failed_count > 0:
        print("\n处理失败的文件:")
        for file in failed_files:
            print(f"  - {file}")
    
    logger.info("所有文件处理完成")
    